/*
Copyright 2023.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package v1

import (
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"

	"github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1/common"
	usagedbapi "github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache/usagedb/api"
)

const (
	binpackStrategy = "binpack"
)

// NOTE: json tags are required.  Any new fields you add must have json tags for the fields to be serialized.

// SchedulingShardSpec defines the desired state of SchedulingShard
type SchedulingShardSpec struct {
	// Args specifies custom CLI arguments for the scheduler. These are merged with automatically generated flags.
	// Valid flags are those defined in the scheduler's code. Usage examples:
	// - To pass "--custom-flag=value": Args: {"custom-flag": "value"}
	// - To override default flags like "--leader-elect": Args: {"leader-elect": "true"}
	// Note:
	//   * Flag names must match exactly (case-sensitive)
	//   * Values are passed as strings
	//   * Only valid flags defined in the scheduler's flag set will be accepted
	//   * Duplicated flags will override the behavior of flags generated by other fields
	// +kubebuilder:validation:Optional
	Args map[string]string `json:"args,omitempty"`

	// PlacementStrategy is the placement scheduler strategy
	// +kubebuilder:validation:Optional
	PlacementStrategy *PlacementStrategy `json:"placementStrategy,omitempty"`

	// PartitionLabelValue is the value for the partition label
	// +kubebuilder:validation:Optional
	PartitionLabelValue string `json:"partitionLabelValue,omitempty"`

	// QueueDepthPerAction max number of jobs to try for action per queue
	// +kubebuilder:validation:Optional
	QueueDepthPerAction map[string]int `json:"queueDepthPerAction,omitempty"`

	// MinRuntime specifies the minimum runtime of a jobs in the shard
	// +kubebuilder:validation:Optional
	MinRuntime *MinRuntime `json:"minRuntime,omitempty"`

	// KValue specifies the kValue for the proportion plugin. Default is 1.0.
	// +kubebuilder:validation:Optional
	KValue *float64 `json:"kValue,omitempty"`

	// UsageDBConfig defines configuration for the usage db client
	// +kubebuilder:validation:Optional
	UsageDBConfig *usagedbapi.UsageDBConfig `yaml:"usageDBConfig,omitempty" json:"usageDBConfig,omitempty"`
}

func (s *SchedulingShardSpec) SetDefaultsWhereNeeded() {
	s.PlacementStrategy = common.SetDefault(s.PlacementStrategy, &PlacementStrategy{})
	s.PlacementStrategy.SetDefaultWhereNeeded()
}

type MinRuntime struct {
	// PreemptMinRuntime specifies the minimum runtime of a job in queue before it can be preempted
	// +kubebuilder:validation:Optional
	PreemptMinRuntime *string `json:"preemptMinRuntime,omitempty"`

	// ReclaimMinRuntime specifies the minimum runtime of a job in queue before it can be reclaimed
	// +kubebuilder:validation:Optional
	ReclaimMinRuntime *string `json:"reclaimMinRuntime,omitempty"`
}

// PlacementStrategy defines the scheduling strategy of NodePool
type PlacementStrategy struct {
	// GPU scheduling strategy (binpack/spread)
	// +kubebuilder:validation:Optional
	GPU *string `json:"gpu,omitempty"`

	// CPU scheduling strategy (binpack/spread)
	// +kubebuilder:validation:Optional
	CPU *string `json:"cpu,omitempty"`
}

func (p *PlacementStrategy) SetDefaultWhereNeeded() {
	p.GPU = common.SetDefault(p.GPU, ptr.To(binpackStrategy))
	p.CPU = common.SetDefault(p.CPU, ptr.To(binpackStrategy))
}

// SchedulingShardStatus defines the observed state of SchedulingShard
type SchedulingShardStatus struct {
	Conditions []metav1.Condition `json:"conditions,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:resource:scope=Cluster

// SchedulingShard is the Schema for the schedulingshards API
type SchedulingShard struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	Spec   SchedulingShardSpec   `json:"spec,omitempty"`
	Status SchedulingShardStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// SchedulingShardList contains a list of SchedulingShard
type SchedulingShardList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []SchedulingShard `json:"items"`
}

func init() {
	SchemeBuilder.Register(&SchedulingShard{}, &SchedulingShardList{})
}
